#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Add retry logic for 504 Gateway Timeout errors in DemandServiceImpl.java
"""

import shutil
import sys

if sys.platform == 'win32':
    import io
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')

file_path = r"C:\Users\lyc\Desktop\spb\src\main\java\com\expert\miniapp\service\impl\DemandServiceImpl.java"

# Create backup
backup_path = file_path + ".backup2"
shutil.copy2(file_path, backup_path)
print(f"[OK] Created backup: {backup_path}")

# Read the file
with open(file_path, 'r', encoding='utf-8') as f:
    content = f.read()

# Find and replace the AI call section (around line 153-212)
old_code = """        HttpClient client = HttpClient.newBuilder()
                .connectTimeout(Duration.ofMillis(Optional.ofNullable(aiConfig.getTimeout()).orElse(180000)))
                .build();
        HttpRequest request = HttpRequest.newBuilder()
                .uri(URI.create(aiConfig.getUrl()))
                .timeout(Duration.ofMillis(Optional.ofNullable(aiConfig.getTimeout()).orElse(180000)))
                .header("Content-Type", "application/json; charset=UTF-8")
                .header("Accept", "application/json")
                .header("Authorization", "Bearer " + aiConfig.getKey())
                .POST(HttpRequest.BodyPublishers.ofString(payloadJson, UTF_8))
                .build();

        String content;
        String responseBody = "";
        try {
            logger.info("=== 开始调用AI匹配接口 ===");
            logger.info("请求URL: {}", aiConfig.getUrl());
            logger.info("请求payload长度: {} 字符", payloadJson.length());

            HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString(UTF_8));
            responseBody = response.body();

            logger.info("=== AI响应诊断信息 ===");
            logger.info("HTTP状态码: {}", response.statusCode());
            logger.info("响应body长度: {} 字符", responseBody != null ? responseBody.length() : 0);
            logger.info("响应body前500字符: {}", responseBody != null && responseBody.length() > 0 ?
                responseBody.substring(0, Math.min(500, responseBody.length())) : "空响应");

            if (response.statusCode() < 200 || response.statusCode() >= 300) {
                logger.error("AI接口返回非2xx状态码: {}, 响应body: {}", response.statusCode(), responseBody);
                throw new RuntimeException("AI匹配接口返回非2xx状态: " + response.statusCode());
            }

            var root = objectMapper.readTree(responseBody);
            logger.info("JSON解析成功");

            var choices = root.path("choices");
            logger.info("choices节点是否存在: {}", !choices.isMissingNode());
            logger.info("choices是否为数组: {}", choices.isArray());
            logger.info("choices数组大小: {}", choices.isArray() ? choices.size() : 0);

            if (!choices.isArray() || choices.isEmpty()) {
                logger.error("AI响应中choices字段无效, 完整响应: {}", responseBody);
                throw new RuntimeException("AI匹配未取得推荐请求采用数据");
            }

            content = choices.get(0).path("message").path("content").asText("");
            logger.info("提取到的content长度: {} 字符", content.length());
            logger.info("content前200字符: {}", content.length() > 0 ?
                content.substring(0, Math.min(200, content.length())) : "空内容");

        } catch (Exception e) {
            logger.error("AI匹配调用失败: {}", e.getMessage(), e);
            logger.error("失败时的响应body: {}", responseBody);
            throw new RuntimeException("AI匹配调用失败: " + e.getMessage(), e);
        }

        if (!StringUtils.hasText(content)) {
            logger.error("AI返回的content为空字符串! 完整响应body: {}", responseBody);
            throw new RuntimeException("AI匹配未返回内容");
        }"""

new_code = """        HttpClient client = HttpClient.newBuilder()
                .connectTimeout(Duration.ofMillis(Optional.ofNullable(aiConfig.getTimeout()).orElse(180000)))
                .build();
        HttpRequest request = HttpRequest.newBuilder()
                .uri(URI.create(aiConfig.getUrl()))
                .timeout(Duration.ofMillis(Optional.ofNullable(aiConfig.getTimeout()).orElse(180000)))
                .header("Content-Type", "application/json; charset=UTF-8")
                .header("Accept", "application/json")
                .header("Authorization", "Bearer " + aiConfig.getKey())
                .POST(HttpRequest.BodyPublishers.ofString(payloadJson, UTF_8))
                .build();

        String content = null;
        String responseBody = "";

        // 重试逻辑：对于504超时错误，最多重试2次
        int maxRetries = 3;
        int retryCount = 0;
        Exception lastException = null;

        while (retryCount < maxRetries && content == null) {
            try {
                if (retryCount > 0) {
                    logger.warn("AI匹配第 {} 次重试 (总共 {} 次)", retryCount + 1, maxRetries);
                    Thread.sleep(5000); // 重试前等待5秒
                }

                logger.info("=== 开始调用AI匹配接口 (尝试 {}/{}) ===", retryCount + 1, maxRetries);
                logger.info("请求URL: {}", aiConfig.getUrl());
                logger.info("请求payload长度: {} 字符", payloadJson.length());

                long startTime = System.currentTimeMillis();
                HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString(UTF_8));
                long elapsed = System.currentTimeMillis() - startTime;
                responseBody = response.body();

                logger.info("=== AI响应诊断信息 ===");
                logger.info("HTTP状态码: {}", response.statusCode());
                logger.info("响应时间: {} 毫秒", elapsed);
                logger.info("响应body长度: {} 字符", responseBody != null ? responseBody.length() : 0);
                logger.info("响应body前500字符: {}", responseBody != null && responseBody.length() > 0 ?
                    responseBody.substring(0, Math.min(500, responseBody.length())) : "空响应");

                // 处理504超时错误 - 可以重试
                if (response.statusCode() == 504) {
                    logger.warn("收到504网关超时错误，可能是AI处理时间过长");
                    retryCount++;
                    lastException = new RuntimeException("AI网关超时(504)，已耗时 " + elapsed + " 毫秒");
                    continue;
                }

                // 其他非2xx错误 - 不重试
                if (response.statusCode() < 200 || response.statusCode() >= 300) {
                    logger.error("AI接口返回非2xx状态码: {}, 响应body: {}", response.statusCode(), responseBody);
                    throw new RuntimeException("AI匹配接口返回非2xx状态: " + response.statusCode());
                }

                var root = objectMapper.readTree(responseBody);
                logger.info("JSON解析成功");

                var choices = root.path("choices");
                logger.info("choices节点是否存在: {}", !choices.isMissingNode());
                logger.info("choices是否为数组: {}", choices.isArray());
                logger.info("choices数组大小: {}", choices.isArray() ? choices.size() : 0);

                if (!choices.isArray() || choices.isEmpty()) {
                    logger.error("AI响应中choices字段无效, 完整响应: {}", responseBody);
                    throw new RuntimeException("AI匹配未取得推荐请求采用数据");
                }

                String tempContent = choices.get(0).path("message").path("content").asText("");
                logger.info("提取到的content长度: {} 字符", tempContent.length());
                logger.info("content前200字符: {}", tempContent.length() > 0 ?
                    tempContent.substring(0, Math.min(200, tempContent.length())) : "空内容");

                if (StringUtils.hasText(tempContent)) {
                    content = tempContent;
                    logger.info("AI匹配成功！第 {} 次尝试获得有效响应", retryCount + 1);
                } else {
                    logger.warn("AI返回的content为空字符串");
                    retryCount++;
                    lastException = new RuntimeException("AI返回的content为空");
                }

            } catch (InterruptedException ie) {
                Thread.currentThread().interrupt();
                throw new RuntimeException("重试等待被中断", ie);
            } catch (Exception e) {
                logger.error("AI匹配调用失败 (尝试 {}/{}): {}", retryCount + 1, maxRetries, e.getMessage(), e);
                logger.error("失败时的响应body: {}", responseBody);
                lastException = e;
                retryCount++;

                // 如果是最后一次尝试，抛出异常
                if (retryCount >= maxRetries) {
                    break;
                }
            }
        }

        // 检查是否成功获取content
        if (!StringUtils.hasText(content)) {
            logger.error("AI匹配失败：经过 {} 次尝试仍未获得有效响应", maxRetries);
            logger.error("完整响应body: {}", responseBody);
            String errorMsg = lastException != null ? lastException.getMessage() : "未知错误";
            throw new RuntimeException("AI匹配未返回内容（已重试" + maxRetries + "次）: " + errorMsg, lastException);
        }"""

# Replace the code
if old_code in content:
    content = content.replace(old_code, new_code)
    print("[OK] Found and replaced AI call section")
else:
    print("[ERROR] Could not find the exact code section to replace")
    print("[INFO] This might mean the file has already been modified or the code structure changed")
    sys.exit(1)

# Write back
with open(file_path, 'w', encoding='utf-8') as f:
    f.write(content)

print(f"[OK] Modification completed!")
print(f"[OK] Added retry logic for 504 Gateway Timeout errors")
print(f"[OK] Original file backed up to: {backup_path}")
print("\nChanges made:")
print("- Added retry loop with max 3 attempts")
print("- Special handling for 504 Gateway Timeout errors")
print("- 5-second delay between retries")
print("- Detailed logging for each attempt")
print("\nPlease restart Spring Boot application for changes to take effect")
